fix(loss): count num_valid_samples from the sample mask, not the batch dim - #3850

…h dim

Six losses in this file derive num_valid_samples from the sample mask.
Three derive it from the batch dimension instead: MseValueLossFn from
values.shape[0], and both distillation losses from input_ids.shape[0].

That number is a gate, not just a log line. dtensor_policy_worker,
dtensor_policy_worker_v2 and dtensor_value_worker_v2 all do

    if num_valid_samples > 0:
        mb_losses.append(loss.item())
        all_mb_metrics.append(loss_metrics)

so a microbatch whose sample_mask is entirely zero -- which is what
overlong filtering, env-flagged masking and the critic's
seq_logprob_error_threshold filter produce -- reports the raw batch size,
passes the gate, and contributes its zero loss to the step's reported mean.
dpo.py and rm.py additionally use the value as a weighted-average
denominator.

DistillationLossFn keeps a batch-dimension fallback: its masking branch is
conditional, and with no mask every sample is valid by definition.

Tests are CPU-only -- they build the loss inputs directly rather than going
through prepare_loss_input, so no GPU and no distributed init. The
cross-tokenizer loss gets the same one-line change but no new test: every
CPU test for it calls the private helpers, and none reaches __call__ where
the metrics dict is built.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

…IA-NeMo#3496

NVIDIA-NeMo#3496 replaces DistillationLossFn's unmasked-mean fallback with a raise. Until
it lands the fallback is live and must report the batch size; after it lands
the branch is gone and the raise is the only correct behaviour.

Asserting one of them unconditionally makes the two PRs fail as a pair while
each passes alone. That is not hypothetical -- it is what happened when I
merged the whole stack onto main to check exactly this.

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>